import sys
import os
import shutil
import cv2
import numpy as np
from pathlib import Path
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning, message="sipPyTypeDict.*")
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                            QPushButton, QLabel, QFileDialog, QProgressBar,
                            QRadioButton, QButtonGroup, QMessageBox, QListWidget,
                            QStackedLayout)
from converter import YoloVocConverter

class ConverterApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.converter = YoloVocConverter()
        self.dataset_path = None
        self.clean_dataset_path = None

    def show_conversion_ui(self):
        """显示转换功能界面"""
        self.stacked_layout.setCurrentIndex(0)

    def show_clean_ui(self):
        """显示文件整理功能界面"""
        self.stacked_layout.setCurrentIndex(1)

    def select_clean_dataset(self):
        """选择要整理的数据集目录"""
        dir_path = QFileDialog.getExistingDirectory(self, '选择数据集目录')
        if dir_path:
            try:
                self.clean_dataset_path = Path(dir_path)
                self.refresh_file_list()
            except Exception as e:
                QMessageBox.warning(self, '错误', f'路径处理失败: {str(e)}')

    def refresh_file_list(self):
        """刷新文件列表"""
        if not self.clean_dataset_path:
            return
            
        self.file_list.clear()
        
        # 根据选择的格式确定要扫描的目录
        if self.clean_yolo_radio.isChecked():
            img_dir = self.clean_dataset_path / 'images'
            label_dir = self.clean_dataset_path / 'labels'
        else:
            img_dir = self.clean_dataset_path / 'JPEGImages'
            label_dir = self.clean_dataset_path / 'Annotations'
            
        # 获取所有图片和标签文件
        img_files = set(f.stem for f in img_dir.glob('*') if f.suffix.lower() in ['.jpg', '.png', '.jpeg'])
        label_files = set(f.stem for f in label_dir.glob('*') if f.suffix.lower() in ['.txt', '.xml'])
        
        # 找出不匹配的文件
        extra_images = img_files - label_files
        extra_labels = label_files - img_files
        
        # 添加到列表
        for img in extra_images:
            self.file_list.addItem(f"多余图片: {img}")
        for label in extra_labels:
            self.file_list.addItem(f"多余标签: {label}")

    def delete_selected_files(self):
        """删除选中的文件"""
        selected_items = [item.text() for item in self.file_list.selectedItems()]
        if not selected_items:
            return
            
        # 确认删除
        reply = QMessageBox.question(self, '确认', 
                                   '确定要删除选中的文件吗？此操作不可恢复！',
                                   QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
            
        # 根据选择的格式确定目录
        if self.clean_yolo_radio.isChecked():
            img_dir = self.clean_dataset_path / 'images'
            label_dir = self.clean_dataset_path / 'labels'
            img_ext = ['.jpg', '.png', '.jpeg']
            label_ext = ['.txt']
        else:
            img_dir = self.clean_dataset_path / 'JPEGImages'
            label_dir = self.clean_dataset_path / 'Annotations'
            img_ext = ['.jpg', '.png', '.jpeg']
            label_ext = ['.xml']
            
        deleted_count = 0
        for item in selected_items:
            try:
                if item.startswith('多余图片:'):
                    filename = item.replace('多余图片:', '').strip()
                    for ext in img_ext:
                        file_path = img_dir / f"{filename}{ext}"
                        if file_path.exists():
                            file_path.unlink()
                            deleted_count += 1
                elif item.startswith('多余标签:'):
                    filename = item.replace('多余标签:', '').strip()
                    for ext in label_ext:
                        file_path = label_dir / f"{filename}{ext}"
                        if file_path.exists():
                            file_path.unlink()
                            deleted_count += 1
            except Exception as e:
                QMessageBox.warning(self, '错误', f'删除文件失败: {str(e)}')
                continue
                
        QMessageBox.information(self, '完成', f'成功删除 {deleted_count} 个文件')
        self.refresh_file_list()

    def initUI(self):
        self.setWindowTitle('YOLO-VOC 数据集工具')
        self.setGeometry(300, 300, 500, 400)
        
        # 主布局
        main_layout = QVBoxLayout()
        
        # 添加数据集路径显示标签
        self.dataset_label = QLabel('数据集目录: 未选择')
        main_layout.addWidget(self.dataset_label)
        
        # 功能选择
        self.function_label = QLabel('请选择功能:')
        main_layout.addWidget(self.function_label)
        
        # 功能按钮
        btn_layout = QHBoxLayout()
        self.convert_btn = QPushButton('数据集转换')
        self.convert_btn.clicked.connect(self.show_conversion_ui)
        btn_layout.addWidget(self.convert_btn)
        
        self.clean_btn = QPushButton('数据文件整理')
        self.clean_btn.clicked.connect(self.show_clean_ui)
        btn_layout.addWidget(self.clean_btn)
        main_layout.addLayout(btn_layout)
        
        # 堆叠布局用于切换不同功能的UI
        self.stacked_layout = QStackedLayout()
        main_layout.addLayout(self.stacked_layout)
        
        # 转换功能UI
        self.conversion_ui = self.create_conversion_ui()
        self.stacked_layout.addWidget(self.conversion_ui)
        
        # 文件整理功能UI
        self.clean_ui = self.create_clean_ui()
        self.stacked_layout.addWidget(self.clean_ui)
        
        self.setLayout(main_layout)
        
    def create_conversion_ui(self):
        """创建转换功能UI"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        # 转换方向选择
        self.direction_label = QLabel('选择转换方向:')
        layout.addWidget(self.direction_label)
        
        self.yolo_to_voc_radio = QRadioButton('YOLO → VOC')
        self.voc_to_yolo_radio = QRadioButton('VOC → YOLO')
        self.format_group = QButtonGroup()
        self.format_group.addButton(self.yolo_to_voc_radio)
        self.format_group.addButton(self.voc_to_yolo_radio)
        
        direction_layout = QHBoxLayout()
        direction_layout.addWidget(self.yolo_to_voc_radio)
        direction_layout.addWidget(self.voc_to_yolo_radio)
        layout.addLayout(direction_layout)
        
        # 数据集路径选择
        self.convert_dataset_label = QLabel('请选择数据集目录:')
        layout.addWidget(self.convert_dataset_label)
        
        self.convert_dataset_btn = QPushButton('选择数据集目录')
        self.convert_dataset_btn.clicked.connect(self.select_dataset)
        layout.addWidget(self.convert_dataset_btn)
        
        # 进度条
        self.progress = QProgressBar()
        layout.addWidget(self.progress)
        
        # 开始转换按钮
        self.start_convert_btn = QPushButton('开始转换')
        self.start_convert_btn.clicked.connect(self.start_conversion)
        layout.addWidget(self.start_convert_btn)
        
        widget.setLayout(layout)
        return widget
        
    def create_clean_ui(self):
        """创建文件整理功能UI"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        # 当前格式选择
        self.clean_format_label = QLabel('选择当前数据集格式:')
        layout.addWidget(self.clean_format_label)
        
        self.clean_yolo_radio = QRadioButton('YOLO')
        self.clean_voc_radio = QRadioButton('VOC')
        self.clean_format_group = QButtonGroup()
        self.clean_format_group.addButton(self.clean_yolo_radio)
        self.clean_format_group.addButton(self.clean_voc_radio)
        
        format_layout = QHBoxLayout()
        format_layout.addWidget(self.clean_yolo_radio)
        format_layout.addWidget(self.clean_voc_radio)
        layout.addLayout(format_layout)
        
        # 数据集路径选择
        self.clean_dataset_label = QLabel('请选择数据集目录:')
        layout.addWidget(self.clean_dataset_label)
        
        self.clean_dataset_btn = QPushButton('选择数据集目录')
        self.clean_dataset_btn.clicked.connect(self.select_clean_dataset)
        layout.addWidget(self.clean_dataset_btn)
        
        # 文件列表
        self.file_list = QListWidget()
        self.file_list.setSelectionMode(QListWidget.MultiSelection)  # 支持多选
        layout.addWidget(self.file_list)
        
        # 操作按钮
        btn_layout = QHBoxLayout()
        self.refresh_btn = QPushButton('刷新列表')
        self.refresh_btn.clicked.connect(self.refresh_file_list)
        btn_layout.addWidget(self.refresh_btn)
        
        self.delete_btn = QPushButton('删除选中文件')
        self.delete_btn.clicked.connect(self.delete_selected_files)
        btn_layout.addWidget(self.delete_btn)
        layout.addLayout(btn_layout)
        
        widget.setLayout(layout)
        return widget
        
    def select_dataset(self):
        """选择数据集目录"""
        dir_path = QFileDialog.getExistingDirectory(self, '选择数据集目录')
        if dir_path:
            try:
                # 处理中文路径
                self.dataset_path = Path(dir_path)
                self.dataset_label.setText(f'数据集目录: {self.dataset_path}')
                
                # 测试路径访问
                test_file = next(self.dataset_path.glob('*'), None)
                if test_file:
                    try:
                        with open(test_file, 'rb') as f:
                            f.read(1)
                    except Exception as e:
                        QMessageBox.warning(self, '错误', f'路径访问失败: {str(e)}')
                        return
            except Exception as e:
                QMessageBox.warning(self, '错误', f'路径处理失败: {str(e)}')
                return
            
    def start_conversion(self):
        """开始转换"""
        if not hasattr(self, 'dataset_path'):
            QMessageBox.warning(self, '错误', '请先选择数据集目录')
            return
            
        if not self.format_group.checkedButton():
            QMessageBox.warning(self, '错误', '请选择当前格式')
            return
            
        # 获取当前格式
        checked_button = self.format_group.checkedButton()
        if checked_button == self.yolo_to_voc_radio:
            is_yolo = True
        else:
            is_yolo = False
        
        try:
            # 执行转换
            if is_yolo:
                self.convert_yolo_to_voc()
            else:
                self.convert_voc_to_yolo()
                
            QMessageBox.information(self, '完成', '转换完成！')
        except Exception as e:
            QMessageBox.critical(self, '错误', f'转换失败: {str(e)}')
            
    def convert_yolo_to_voc(self):
        """YOLO转VOC"""
        # 创建输出目录
        output_dir = self.dataset_path / 'VOC'
        output_dir.mkdir(exist_ok=True)
        
        # 创建VOC标准目录结构
        jpeg_dir = output_dir / 'JPEGImages'
        annot_dir = output_dir / 'Annotations'
        jpeg_dir.mkdir(exist_ok=True)
        annot_dir.mkdir(exist_ok=True)
        
        # 获取类别列表
        classes_file = self.dataset_path / 'classes.txt'
        if not classes_file.exists():
            raise FileNotFoundError('未找到classes.txt文件')
            
        with open(classes_file, 'r', encoding='utf-8') as f:
            classes = [line.strip() for line in f.readlines()]
        self.converter.set_classes(classes)
        
        # 获取所有图片文件
        img_dir = self.dataset_path / 'images'
        if not img_dir.exists():
            raise FileNotFoundError('未找到images目录')
            
        img_files = list(img_dir.glob('*.jpg')) + list(img_dir.glob('*.png'))
        self.progress.setMaximum(len(img_files))
        
        # 处理每张图片
        for i, img_path in enumerate(img_files):
            # 更新进度
            self.progress.setValue(i + 1)
            QApplication.processEvents()
            
            # 获取对应的标注文件
            label_path = self.dataset_path / 'labels' / f"{img_path.stem}.txt"
            if not label_path.exists():
                continue
                
            # 检查图片文件是否存在
            if not img_path.exists():
                raise FileNotFoundError(f'图片文件不存在: {img_path}')
                
            # 读取图片尺寸
            try:
                # 使用PIL打开图片，支持中文路径
                from PIL import Image
                pil_img = Image.open(str(img_path))
                # 转换为OpenCV支持的格式
                img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            except Exception as e:
                print(f'警告：无法读取图片文件 {img_path}，跳过该文件。错误详情: {str(e)}')
                continue
                    
            img_height, img_width = img.shape[:2]
            
            # 读取YOLO标注
            with open(label_path, 'r', encoding='utf-8') as f:
                yolo_txt = f.read()
                
            # 转换为VOC格式
            objects = self.converter.yolo_to_voc(yolo_txt, img_width, img_height)
            
            # 保存VOC XML
            output_img_path = jpeg_dir / img_path.name
            output_xml_path = annot_dir / f"{img_path.stem}.xml"
            
            # 复制图片
            shutil.copy(img_path, output_img_path)
            
            # 保存XML
            self.converter.save_voc_xml(
                output_img_path,
                (img_width, img_height),
                objects,
                output_xml_path
            )
            
    def find_extra_files(self):
        """查找多余的文件"""
        if not hasattr(self, 'dataset_path'):
            QMessageBox.warning(self, '错误', '请先选择数据集目录')
            return None, None
            
        # 根据当前格式确定目录结构
        if self.yolo_radio.isChecked():
            img_dir = self.dataset_path / 'images'
            label_dir = self.dataset_path / 'labels'
            img_ext = ['.jpg', '.png', '.jpeg']
            label_ext = ['.txt']
        else:
            img_dir = self.dataset_path / 'JPEGImages'
            label_dir = self.dataset_path / 'Annotations'
            img_ext = ['.jpg', '.png', '.jpeg']
            label_ext = ['.xml']
            
        # 检查目录是否存在
        if not img_dir.exists():
            QMessageBox.warning(self, '错误', f'图片目录不存在: {img_dir}')
            return None, None
        if not label_dir.exists():
            QMessageBox.warning(self, '错误', f'标签目录不存在: {label_dir}')
            return None, None
            
        # 获取所有图片和标签文件
        img_files = set(f.stem for f in img_dir.glob('*') if f.suffix.lower() in img_ext)
        label_files = set(f.stem for f in label_dir.glob('*') if f.suffix.lower() in label_ext)
        
        # 找出多余的文件
        extra_images = img_files - label_files
        extra_labels = label_files - img_files
        
        return extra_images, extra_labels
        
    def clean_extra_files(self):
        """清理多余文件"""
        extra_images, extra_labels = self.find_extra_files()
        if extra_images is None:
            return
            
        # 清空文件列表
        self.file_list.clear()
        
        # 添加多余图片到列表
        for img in extra_images:
            self.file_list.addItem(f"多余图片: {img}")
            
        # 添加多余标签到列表
        for label in extra_labels:
            self.file_list.addItem(f"多余标签: {label}")
            
        if not extra_images and not extra_labels:
            QMessageBox.information(self, '提示', '没有发现多余文件')
            return
            
        # 弹出确认对话框
        reply = QMessageBox.question(self, '确认', 
                                   '是否要删除选中的多余文件？',
                                   QMessageBox.Yes | QMessageBox.No)
        
        if reply == QMessageBox.Yes:
            # 获取选中的项目
            selected_items = [item.text() for item in self.file_list.selectedItems()]
            
            # 根据当前格式确定目录结构
            if self.yolo_radio.isChecked():
                img_dir = self.dataset_path / 'images'
                label_dir = self.dataset_path / 'labels'
                img_ext = ['.jpg', '.png', '.jpeg']
                label_ext = ['.txt']
            else:
                img_dir = self.dataset_path / 'JPEGImages'
                label_dir = self.dataset_path / 'Annotations'
                img_ext = ['.jpg', '.png', '.jpeg']
                label_ext = ['.xml']
            
            # 删除选中的文件
            deleted_count = 0
            for item in selected_items:
                try:
                    if item.startswith('多余图片:'):
                        filename = item.replace('多余图片:', '').strip()
                        for ext in img_ext:
                            file_path = img_dir / f"{filename}{ext}"
                            if file_path.exists():
                                file_path.unlink()
                                deleted_count += 1
                    elif item.startswith('多余标签:'):
                        filename = item.replace('多余标签:', '').strip()
                        for ext in label_ext:
                            file_path = label_dir / f"{filename}{ext}"
                            if file_path.exists():
                                file_path.unlink()
                                deleted_count += 1
                except Exception as e:
                    QMessageBox.warning(self, '错误', f'删除文件失败: {str(e)}')
                    continue
                    
            QMessageBox.information(self, '完成', f'成功删除 {deleted_count} 个文件')
            
    def convert_voc_to_yolo(self):
        """VOC转YOLO"""
        # 创建输出目录
        output_dir = self.dataset_path / 'YOLO'
        output_dir.mkdir(exist_ok=True)
        
        # 创建YOLO标准目录结构
        img_dir = output_dir / 'images'
        label_dir = output_dir / 'labels'
        img_dir.mkdir(exist_ok=True)
        label_dir.mkdir(exist_ok=True)
        
        # 获取类别列表
        classes_file = self.dataset_path / 'classes.txt'
        if not classes_file.exists():
            raise FileNotFoundError('未找到classes.txt文件')
            
        with open(classes_file, 'r', encoding='utf-8') as f:
            classes = [line.strip() for line in f.readlines()]
        self.converter.set_classes(classes)
        
        # 获取所有图片文件
        jpeg_dir = self.dataset_path / 'JPEGImages'
        if not jpeg_dir.exists():
            raise FileNotFoundError('未找到JPEGImages目录')
            
        img_files = list(jpeg_dir.glob('*.jpg')) + list(jpeg_dir.glob('*.png'))
        self.progress.setMaximum(len(img_files))
        
        # 处理每张图片
        for i, img_path in enumerate(img_files):
            # 更新进度
            self.progress.setValue(i + 1)
            QApplication.processEvents()
            
            # 获取对应的标注文件
            xml_path = self.dataset_path / 'Annotations' / f"{img_path.stem}.xml"
            if not xml_path.exists():
                continue
                
            # 读取图片尺寸
            try:
                # 使用PIL打开图片，支持中文路径
                from PIL import Image
                pil_img = Image.open(str(img_path))
                # 转换为OpenCV支持的格式
                img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
                img_height, img_width = img.shape[:2]
            except Exception as e:
                print(f'警告：无法读取图片文件 {img_path}，跳过该文件。错误详情: {str(e)}')
                continue
            
            # 读取VOC标注
            voc_annots = self.converter.parse_voc_xml(xml_path)
            
            # 转换为YOLO格式
            yolo_lines = self.converter.voc_to_yolo(voc_annots, img_width, img_height)
            
            # 保存YOLO标注
            output_img_path = img_dir / img_path.name
            output_label_path = label_dir / f"{img_path.stem}.txt"
            
            # 复制图片
            shutil.copy(img_path, output_img_path)
            
            # 保存YOLO标注
            with open(output_label_path, 'w', encoding='utf-8') as f:
                f.write('\n'.join(yolo_lines))

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = ConverterApp()
    window.show()
    sys.exit(app.exec_())
